#ifndef UNICODE
#define UNICODE
#endif

#include <windows.h>
#include <dbt.h>
#include <setupapi.h>
#include <stdio.h>
#include <initguid.h>
#include <devguid.h>

#pragma comment(lib, "setupapi.lib")

// 定义HID GUID (缺失的定义)
DEFINE_GUID(GUID_DEVINTERFACE_HID, 0x4D1E55B2L, 0xF16F, 0x11CF, 0x88, 0xCB, 0x00, 0x11, 0x11, 0x00, 0x00, 0x30);

// 定义窗口类名
#define WINDOW_CLASS_NAME TEXT("DeviceModeMonitorClass")

// 支持的目标设备硬件ID
#define MAX_TARGET_DEVICES 2
const WCHAR* TARGET_HARDWARE_IDS[MAX_TARGET_DEVICES] = {
    L"HID#VID_045E&PID_028E",
    L"HID#VID_258A&PID_001E&MI_00#9&B00F61D&0&0000"
};

// 模式定义
#define DESKTOP_MODE 1
#define TABLET_MODE 0

// 服务名称
#define SERVICE_NAME TEXT("DeviceModeMonitor")
#define SERVICE_DISPLAY_NAME TEXT("Device Mode Monitor")
#define SERVICE_DESCRIPTION TEXT("Monitors devices and switches between desktop and tablet modes")

// 添加全局变量跟踪当前模式状态
// -1表示初始未知状态，0表示平板模式，1表示桌面模式
int currentMode = -1;

// 服务状态和控制句柄
SERVICE_STATUS gServiceStatus = {0};
SERVICE_STATUS_HANDLE gServiceStatusHandle = NULL;
HANDLE gServiceStopEvent = NULL;

// 函数声明
LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam);
void SwitchToDesktopMode();
void SwitchToTabletMode();
BOOL IsAnyTargetDevicePresent();
void CheckAndUpdateMode();
BOOL IsTargetDevice(LPCWSTR devicePath);
int GetCurrentSystemMode();

// 服务相关函数声明
VOID WINAPI ServiceMain(DWORD argc, LPTSTR *argv);
VOID WINAPI ServiceCtrlHandler(DWORD controlCode);
BOOL InitializeService();
BOOL ReportServiceStatus(DWORD currentState, DWORD exitCode, DWORD waitHint);
DWORD WINAPI ServiceWorkerThread(LPVOID lpParam);
int RunDeviceMonitor();
BOOL IsRunningAsService();

// 主函数 - 会根据运行模式分发到服务主函数或直接运行
int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow) {
    // 判断是否作为服务启动
    if (IsRunningAsService()) {
        // 作为服务运行
        SERVICE_TABLE_ENTRY serviceTable[] = {
            { SERVICE_NAME, ServiceMain },
            { NULL, NULL }
        };

        if (!StartServiceCtrlDispatcher(serviceTable)) {
            // 如果StartServiceCtrlDispatcher失败，说明不是由服务控制管理器启动的
            // 这种情况下以普通应用程序模式运行
            return RunDeviceMonitor();
        }
    } else {
        // 作为普通应用程序运行
        return RunDeviceMonitor();
    }

    return 0;
}

// 服务主函数 - 当程序作为服务启动时，这是入口点
VOID WINAPI ServiceMain(DWORD argc, LPTSTR *argv) {
    // 注册服务控制处理函数
    gServiceStatusHandle = RegisterServiceCtrlHandler(SERVICE_NAME, ServiceCtrlHandler);
    if (!gServiceStatusHandle) {
        return;
    }

    // 设置服务正在启动
    gServiceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
    gServiceStatus.dwControlsAccepted = 0;
    gServiceStatus.dwCurrentState = SERVICE_START_PENDING;
    gServiceStatus.dwWin32ExitCode = 0;
    gServiceStatus.dwServiceSpecificExitCode = 0;
    gServiceStatus.dwCheckPoint = 0;
    gServiceStatus.dwWaitHint = 0;

    if (!ReportServiceStatus(SERVICE_START_PENDING, NO_ERROR, 3000)) {
        ReportServiceStatus(SERVICE_STOPPED, GetLastError(), 0);
        return;
    }

    // 创建服务停止事件
    gServiceStopEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (gServiceStopEvent == NULL) {
        ReportServiceStatus(SERVICE_STOPPED, GetLastError(), 0);
        return;
    }

    // 初始化服务
    if (!InitializeService()) {
        ReportServiceStatus(SERVICE_STOPPED, GetLastError(), 0);
        return;
    }

    // 报告服务正在运行
    if (!ReportServiceStatus(SERVICE_RUNNING, NO_ERROR, 0)) {
        ReportServiceStatus(SERVICE_STOPPED, GetLastError(), 0);
        return;
    }

    // 启动服务工作线程
    HANDLE hThread = CreateThread(NULL, 0, ServiceWorkerThread, NULL, 0, NULL);
    if (hThread) {
        // 等待服务停止信号
        WaitForSingleObject(gServiceStopEvent, INFINITE);
        // 等待工作线程结束
        WaitForSingleObject(hThread, 5000);
        CloseHandle(hThread);
    }

    // 清理
    if (gServiceStopEvent) {
        CloseHandle(gServiceStopEvent);
        gServiceStopEvent = NULL;
    }

    ReportServiceStatus(SERVICE_STOPPED, NO_ERROR, 0);
}

// 服务控制处理函数 - 处理服务控制请求
VOID WINAPI ServiceCtrlHandler(DWORD controlCode) {
    switch (controlCode) {
        case SERVICE_CONTROL_STOP:
            if (gServiceStatus.dwCurrentState != SERVICE_RUNNING)
                break;

            // 报告服务停止中
            ReportServiceStatus(SERVICE_STOP_PENDING, NO_ERROR, 0);

            // 发送停止事件信号
            if (gServiceStopEvent)
                SetEvent(gServiceStopEvent);
            break;

        case SERVICE_CONTROL_INTERROGATE:
            // 仅报告当前状态
            break;

        default:
            break;
    }

    ReportServiceStatus(gServiceStatus.dwCurrentState, NO_ERROR, 0);
}

// 初始化服务
BOOL InitializeService() {
    // 此函数用于初始化服务所需资源
    // 对于此应用，大部分初始化在工作线程中进行
    return TRUE;
}

// 报告服务状态
BOOL ReportServiceStatus(DWORD currentState, DWORD exitCode, DWORD waitHint) {
    // 非首次调用的验证
    static DWORD checkPoint = 1;

    // 填充服务状态结构
    gServiceStatus.dwCurrentState = currentState;
    gServiceStatus.dwWin32ExitCode = exitCode;
    gServiceStatus.dwWaitHint = waitHint;

    // 根据服务状态更新接受的控制命令
    if (currentState == SERVICE_START_PENDING) {
        gServiceStatus.dwControlsAccepted = 0;
    } else {
        gServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    }

    // 更新检查点
    if ((currentState == SERVICE_RUNNING) || (currentState == SERVICE_STOPPED)) {
        gServiceStatus.dwCheckPoint = 0;
    } else {
        gServiceStatus.dwCheckPoint = checkPoint++;
    }

    // 报告状态
    return SetServiceStatus(gServiceStatusHandle, &gServiceStatus);
}

// 服务工作线程 - 执行实际的设备监控任务
DWORD WINAPI ServiceWorkerThread(LPVOID lpParam) {
    // 运行设备监控功能
    RunDeviceMonitor();
    return 0;
}

// 判断程序是否以服务方式运行
BOOL IsRunningAsService() {
    // 尝试访问服务特定的环境
    SC_HANDLE scm = OpenSCManager(NULL, NULL, SC_MANAGER_CONNECT);
    if (!scm) {
        // 没有权限访问服务控制管理器，可能不是以服务方式运行
        return FALSE;
    }
    CloseServiceHandle(scm);

    // 检查是否是由服务控制管理器启动的
    HANDLE hProcess = GetCurrentProcess();
    HANDLE hToken = NULL;
    BOOL result = FALSE;

    if (OpenProcessToken(hProcess, TOKEN_QUERY, &hToken)) {
        TOKEN_STATISTICS stats;
        DWORD retSize = 0;

        if (GetTokenInformation(hToken, TokenStatistics, &stats, sizeof(stats), &retSize)) {
            // 服务通常以LocalSystem账户运行
            SID_IDENTIFIER_AUTHORITY ntAuth = SECURITY_NT_AUTHORITY;
            PSID systemSid = NULL;

            if (AllocateAndInitializeSid(&ntAuth, 1, SECURITY_LOCAL_SYSTEM_RID, 0, 0, 0, 0, 0, 0, 0, &systemSid)) {
                BOOL isSystem = FALSE;
                if (CheckTokenMembership(NULL, systemSid, &isSystem) && isSystem) {
                    result = TRUE;
                }
                FreeSid(systemSid);
            }
        }
        CloseHandle(hToken);
    }

    return result;
}

// 设备监控功能的主实现
int RunDeviceMonitor() {
    HINSTANCE hInstance = GetModuleHandle(NULL);

    // 注册窗口类
    WNDCLASSEX wc = {0};
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = WINDOW_CLASS_NAME;

    if (!RegisterClassEx(&wc)) {
        return 1;
    }

    // 创建一个完全隐藏的窗口
    // 使用WS_EX_TOOLWINDOW | WS_EX_TRANSPARENT | WS_EX_NOACTIVATE样式组合确保窗口不会出现
    HWND hwnd = CreateWindowEx(
        WS_EX_TOOLWINDOW | WS_EX_TRANSPARENT | WS_EX_NOACTIVATE,  // 增加NOACTIVATE和TRANSPARENT标志
        WINDOW_CLASS_NAME,
        TEXT(""),  // 空标题
        WS_POPUP,  // 只使用WS_POPUP样式
        0, 0, 0, 0,  // 位置和尺寸都设为0
        NULL, NULL, hInstance, NULL
    );

    if (hwnd == NULL) {
        return 1;
    }

    // 将窗口移到屏幕外 (额外保证不可见)
    SetWindowPos(hwnd, NULL, -100, -100, 0, 0, SWP_NOZORDER | SWP_NOSIZE | SWP_NOACTIVATE);

    // 显示窗口但使用SW_HIDE确保它绝对不可见
    ShowWindow(hwnd, SW_HIDE);
    
    // 注册设备变更通知
    DEV_BROADCAST_DEVICEINTERFACE notificationFilter = {0};
    notificationFilter.dbcc_size = sizeof(DEV_BROADCAST_DEVICEINTERFACE);
    notificationFilter.dbcc_devicetype = DBT_DEVTYP_DEVICEINTERFACE;
    notificationFilter.dbcc_classguid = GUID_DEVINTERFACE_HID;

    HDEVNOTIFY deviceNotificationHandle = RegisterDeviceNotification(
        hwnd,
        &notificationFilter,
        DEVICE_NOTIFY_WINDOW_HANDLE
    );

    if (deviceNotificationHandle == NULL) {
        DestroyWindow(hwnd);
        return 1;
    }

    // 检查当前设备状态并设置初始模式
    CheckAndUpdateMode();

    // 消息循环 - 对于服务模式，需要检查服务停止事件
    MSG msg = {0};
    BOOL bRet;
    while ((bRet = GetMessage(&msg, NULL, 0, 0)) != 0) {
        if (bRet == -1) {
            // 消息获取失败
            break;
        }

        // 如果正在以服务方式运行，检查停止事件
        if (IsRunningAsService() && gServiceStopEvent != NULL) {
            if (WaitForSingleObject(gServiceStopEvent, 0) == WAIT_OBJECT_0) {
                // 收到停止信号，跳出循环
                break;
            }
        }

        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    // 注销设备通知
    UnregisterDeviceNotification(deviceNotificationHandle);
    DestroyWindow(hwnd);

    return 0;
}

LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
    switch (uMsg) {
        case WM_DEVICECHANGE: {
            if (wParam == DBT_DEVICEARRIVAL || wParam == DBT_DEVICEREMOVECOMPLETE) {
                // 获取设备接口信息
                PDEV_BROADCAST_HDR pHdr = (PDEV_BROADCAST_HDR)lParam;
                if (pHdr->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE) {
                    PDEV_BROADCAST_DEVICEINTERFACE pDevInterface = (PDEV_BROADCAST_DEVICEINTERFACE)pHdr;

                    // 检查是否是目标设备
                    if (IsTargetDevice(pDevInterface->dbcc_name)) {
                        // 只有目标设备变化时才更新模式
                        CheckAndUpdateMode();
                    }
                }
            }
            break;
        }
        case WM_DESTROY: {
            PostQuitMessage(0);
            break;
        }
        default:
            return DefWindowProc(hwnd, uMsg, wParam, lParam);
    }
    return 0;
}

// 判断设备路径是否匹配任一目标设备
BOOL IsTargetDevice(LPCWSTR devicePath) {
    if (!devicePath) return FALSE;

    WCHAR upperDevicePath[MAX_PATH];
    wcsncpy_s(upperDevicePath, MAX_PATH, devicePath, _TRUNCATE);
    _wcsupr_s(upperDevicePath, MAX_PATH);

    BOOL result = FALSE;
    for (int i = 0; i < MAX_TARGET_DEVICES; i++) {
        WCHAR upperTargetId[256];
        wcsncpy_s(upperTargetId, 256, TARGET_HARDWARE_IDS[i], _TRUNCATE);
        _wcsupr_s(upperTargetId, 256);

        if (wcsstr(upperDevicePath, upperTargetId) != NULL) {
            result = TRUE;
            break;
        }
    }

    return result;
}

// 检查设备状态并更新模式
void CheckAndUpdateMode() {
    BOOL devicePresent = IsAnyTargetDevicePresent();
    
    if (currentMode == -1) {
        currentMode = GetCurrentSystemMode();
    }
    
    if (devicePresent) {
        SwitchToDesktopMode();
    } else {
        SwitchToTabletMode();
    }
}

// 检查是否有任一目标设备存在
BOOL IsAnyTargetDevicePresent() {
    BOOL deviceFound = FALSE;
    HDEVINFO deviceInfoSet = SetupDiGetClassDevs(
        &GUID_DEVINTERFACE_HID, 
        NULL, 
        NULL, 
        DIGCF_PRESENT | DIGCF_DEVICEINTERFACE
    );

    if (deviceInfoSet == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    SP_DEVICE_INTERFACE_DATA interfaceData = {0};
    interfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);

    for (DWORD i = 0; SetupDiEnumDeviceInterfaces(deviceInfoSet, NULL, &GUID_DEVINTERFACE_HID, i, &interfaceData); i++) {
        DWORD requiredSize = 0;
        SetupDiGetDeviceInterfaceDetail(deviceInfoSet, &interfaceData, NULL, 0, &requiredSize, NULL);

        DWORD error = GetLastError();
        if (error != ERROR_INSUFFICIENT_BUFFER) {
            continue;
        }

        PSP_DEVICE_INTERFACE_DETAIL_DATA detailData = (PSP_DEVICE_INTERFACE_DETAIL_DATA)malloc(requiredSize);
        if (!detailData) {
            continue;
        }

        detailData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);

        SP_DEVINFO_DATA deviceInfoData = {0};
        deviceInfoData.cbSize = sizeof(SP_DEVINFO_DATA);

        if (SetupDiGetDeviceInterfaceDetail(deviceInfoSet, &interfaceData, detailData, requiredSize, NULL, &deviceInfoData)) {
            if (IsTargetDevice(detailData->DevicePath)) {
                deviceFound = TRUE;
                free(detailData);
                break;
            }
        }

        free(detailData);
    }

    SetupDiDestroyDeviceInfoList(deviceInfoSet);

    return deviceFound;
}

// 切换到桌面模式
void SwitchToDesktopMode() {
    int systemMode = GetCurrentSystemMode();
    
    if (systemMode == DESKTOP_MODE && currentMode == DESKTOP_MODE) {
        return;
    }

    HKEY hKey;
    if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, TEXT("SYSTEM\\CurrentControlSet\\Control\\PriorityControl"), 0, KEY_WRITE, &hKey) == ERROR_SUCCESS) {
        DWORD value = 1;
        RegSetValueEx(hKey, TEXT("ConvertibleSlateMode"), 0, REG_DWORD, (BYTE*)&value, sizeof(DWORD));
        RegCloseKey(hKey);
    }
    
    currentMode = DESKTOP_MODE;
}

// 切换到平板模式
void SwitchToTabletMode() {
    int systemMode = GetCurrentSystemMode();
    
    if (systemMode == TABLET_MODE && currentMode == TABLET_MODE) {
        return;
    }

    HKEY hKey;
    if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, TEXT("SYSTEM\\CurrentControlSet\\Control\\PriorityControl"), 0, KEY_WRITE, &hKey) == ERROR_SUCCESS) {
        DWORD value = 0;
        RegSetValueEx(hKey, TEXT("ConvertibleSlateMode"), 0, REG_DWORD, (BYTE*)&value, sizeof(DWORD));
        RegCloseKey(hKey);
    }
    
    currentMode = TABLET_MODE;
}

// 获取当前系统模式
int GetCurrentSystemMode() {
    HKEY hKey;
    DWORD value = 0;
    DWORD dataSize = sizeof(DWORD);
    
    if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, TEXT("SYSTEM\\CurrentControlSet\\Control\\PriorityControl"), 0, KEY_READ, &hKey) == ERROR_SUCCESS) {
        RegQueryValueEx(hKey, TEXT("ConvertibleSlateMode"), NULL, NULL, (BYTE*)&value, &dataSize);
        RegCloseKey(hKey);
    }
    
    return (int)value;
}
